from pdb import set_trace
from PIL import ImageFont, ImageDraw, Image, ImageFilter
import numpy as np
import sys
sys.path = ['/home/hzh/.local/lib/python3.5/site-packages/'] + sys.path
import cv2
import torch


def get_channel(img):
    if len(img.shape) == 2:
        return 1
    else:
        return img.shape[2]

def remove_row_margin(img):
    rows = img.shape[0]
    top, down = 0, rows-1
    for t in range(0, rows):
        s = np.sum(img[t, :] > 0)
        if(s > 3):
            top = t
            break
    for d in range(rows-1, top, -1):
        s = np.sum(img[d, :] > 0)
        if(s > 3):
            down = d
            break
    return img[top:down+1, :]

def remove_col_margin(img):
    cols = img.shape[1]
    left, right = 0, cols-1
    # set_trace()
    for l in range(0, cols):
        s = np.sum(img[:, l] > 0)
        if(s > 3):
            # print('s=%d, l=%s' % (s, l))
            left = l
            break
    else:
        left = cols

    for r in range(cols-1, left, -1):
        s = np.sum(img[:, r] > 0)
        if(s > 3):
            # print('s=%d, r=%s' % (s, r))
            right = r
            break
    else:
        right = 0

    return img[:, left:right+1] if left < right else np.zeros((1, 1), dtype=img.dtype)

class RemoveMargin:
    def tr(self, img):
        # a = np.asarray(img)
        a = remove_col_margin(remove_row_margin(img))
        # print('RemoveMargin ', img.shape, '-->', a.shape)
        return a
        # c = Image.fromarray(b)

    def __call__(self, imgs):
        if type(imgs) != list:
            return self.tr(imgs)
        new_imgs = []
        for im in imgs:
            new_imgs.append(self.tr(im))
        return new_imgs


class UndistortResize:

    def __init__(self, l, channel):
        self.std_sz = (l, l)
    
    def __pad_3channel__(self, img):
        assert img.dtype == np.uint8, "wrong dtype of image:%s"%img.dtype
        rows, cols, ch = img.shape
        if rows < cols:
            dr = cols - rows
            pack = (img, np.zeros((dr, cols, ch), dtype=np.uint8) )
            img = np.vstack(pack)

        if cols < rows:
            dc = rows - cols
            pack = (img, np.zeros((rows, dc, ch), dtype=np.uint8) )
            img = np.hstack(pack)
        assert img.shape[0] == img.shape[1]
        assert img.shape[2] == 3
        return img

    def __pad_1channel__(self, img):
        assert img.dtype == np.uint8, "wrong dtype of image:%s"%img.dtype
        rows, cols = img.shape
        if rows < cols:
            dr = cols - rows
            pack = (img, np.zeros((dr, cols), dtype=np.uint8) )
            img = np.vstack(pack)

        if cols < rows:
            dc = rows - cols
            pack = (img, np.zeros((rows, dc), dtype=np.uint8) )
            img = np.hstack(pack)
        assert img.shape[0] == img.shape[1]
        return img
    
    def __pad__(self, img):
        if len(img.shape) == 3:
            return self.__pad_3channel__(img)
        else:
            return self.__pad_1channel__(img)

    def __call__(self, img):
        img = self.__pad__(img)
        return cv2.resize(img, self.std_sz)
        

    
class ToGray:

    def __call__(self, img):
        if 1 == get_channel(img):
            return img
        else:
            return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

class Scalarify:

    def __call__(self, img):
        img[img > 0] = WHITE
        return img

    
class CV_pipeline:

    def __init__(self, *processor):
        self._process = [p for p in processor]

    def __call__(self, img):
        for p in self._process:
            img = p(img)
        return img
    
class Reshape2Tensor:
    
    def __call__(self, img):
        h, w = img.shape
        return img.reshape((1, h, w))
    

class ExtraPadding:
    
    def __init__(self, h, w):
        self._w = w
        self._h = h
        
    def __call__(self, img):
        cols = img.shape[1]
        extra_rows = np.zeros((self._h, cols, 3), dtype=np.uint8)
        img = np.vstack((extra_rows, img, extra_rows))
        rows = img.shape[0]
        extra_cols = np.zeros((rows, self._w, 3), dtype=np.uint8)
        img = np.hstack((extra_cols, img, extra_cols))
        return img
        

if '__main__' == __name__:
    img = np.zeros((20, 20, 3), dtype=np.uint8)
    e = ExtraPadding(10, 10)
    out = e(img)
    assert out.shape == (40, 40, 3)
    print("unitTest Success")